/**
 * This program is free software, you can redistribute it and/or modify it.
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * This file is a part of the CANN Open Software.
 * Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
 * BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 */

/*!
 * \file reduce_log_sum_exp_tiling.cpp
 * \brief
*/

#include "log/log.h"
#include "util/math_util.h"
#include "tiling_base/tiling_util.h"
#include "tiling/platform/platform_ascendc.h"
#include "register/op_impl_registry.h"
#include "tiling_base/tiling_templates_registry.h"
#include "../op_kernel/reduce_log_sum_exp_tiling_data.h"
#include "../op_kernel/reduce_log_sum_exp_tiling_key.h"

namespace optiling {

using namespace Ops::Math::OpTiling;
using namespace std;

constexpr uint32_t BLOCK_SIZE = 32;    // 32B对齐单位
constexpr uint32_t BUFFER_NUM = 2;     // 双缓冲块数

#define UB_DATA_NUM_FLOAT16 7 // 对应DT_FLOAT16类型的ub分块数量
#define UB_DATA_NUM_FLOAT32 5 // 对应DT_FLOAT类型的ub分块数量
struct ReduceLogSumExpCompileInfo {};

static ge::graphStatus TilingParseForReduceLogSumExp([[maybe_unused]] gert::TilingParseContext* context)
{
    OP_CHECK_IF(context == nullptr, OP_LOGE(context, "context is nullptr"), return ge::GRAPH_FAILED);
    return ge::GRAPH_SUCCESS;
}

static ge::graphStatus GetPlatformInfo(gert::TilingContext* context, uint64_t& ubSize, uint64_t& coreNum, uint32_t& inputTypeLen)
{
    OP_CHECK_IF(context == nullptr, OP_LOGE(context, "context is nullptr"), return ge::GRAPH_FAILED);
    // 获取ubsize coreNum
    auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
    ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize);
    coreNum = ascendcPlatform.GetCoreNum();
    inputTypeLen = 0;
    ge::TypeUtils::GetDataTypeLength(context->GetInputDesc(0)->GetDataType(), inputTypeLen);
    OP_CHECK_IF(coreNum <= 0, OP_LOGE(context, "coreNum is 0"), return ge::GRAPH_FAILED);
    OP_CHECK_IF(ubSize <= 0, OP_LOGE(context, "ubSize is 0"), return ge::GRAPH_FAILED);
    OP_CHECK_IF(context->GetInputDesc(0)->GetDataType() != ge::DT_FLOAT16 && context->GetInputDesc(0)->GetDataType() != ge::DT_FLOAT, OP_LOGE(context, "xType is not supported"), return ge::GRAPH_FAILED);
    return ge::GRAPH_SUCCESS;
}
    
static ge::graphStatus GetWorkspaceSize(gert::TilingContext* context)
{
    OP_CHECK_IF(context == nullptr, OP_LOGE(context, "context is nullptr"), return ge::GRAPH_FAILED);
    size_t usrSize = 0;
    auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
    uint32_t sysWorkspaceSize = ascendcPlatform.GetLibApiWorkSpaceSize();
    size_t* currentWorkspace = context->GetWorkspaceSizes(1); // 通过框架获取workspace的指针，GetWorkspaceSizes入参为所需workspace的块数。当前限制使用一块。
    currentWorkspace[0] = usrSize + sysWorkspaceSize;
    return ge::GRAPH_SUCCESS;
}

static ge::graphStatus GetTilingAttrsInfo(
    gert::TilingContext* context, uint64_t ubSize, uint32_t& inputNum, uint32_t& tileDataNum, uint32_t& finaltileNum, uint32_t& tailDataNum,
    uint32_t& coreDataNum, uint32_t& tailCoreNum, uint32_t& cycleDataNum, uint32_t& cycleNum, uint32_t& axesNum, uint32_t& shapeDimNum, vector<int32_t> &axes, vector<int32_t> &shapeNum, vector<int32_t> &strides, uint64_t coreNum, uint32_t inputTypeLen, vector<bool> &reduceMask)
{
    OP_CHECK_IF(
    context == nullptr || context->GetInputShape(0) == nullptr || context->GetInputShape(1) == nullptr, OP_LOGE(context, "context is nullptr"),
    return ge::GRAPH_FAILED);
    inputNum = context->GetInputShape(0)->GetStorageShape().GetShapeSize(); 
    if (inputNum == 0) {
        return ge::GRAPH_FAILED;
    }       
    axesNum = context->GetInputShape(1)->GetStorageShape().GetShapeSize();
    axes.resize(axesNum, 0);
    uint32_t axesTypeLen = 0;
    ge::TypeUtils::GetDataTypeLength(context->GetInputDesc(1)->GetDataType(), axesTypeLen);
    OP_CHECK_IF(context->GetInputDesc(1)->GetDataType() != ge::DT_INT32 && context->GetInputDesc(1)->GetDataType() != ge::DT_INT64, OP_LOGE(context, "axesType is not supported"), return ge::GRAPH_FAILED);    
    if (context->GetInputDesc(1)->GetDataType() == ge::DT_INT32) {
        for (uint32_t i = 0; i < axesNum; i++) {
            axes[i] = (context->GetInputTensor(1)->GetData<int32_t>())[i];
        }
    }
    else {
        for (uint32_t i = 0; i < axesNum; i++) {
            axes[i] = (context->GetInputTensor(1)->GetData<int64_t>())[i];
        }
    }
    std::sort(axes.begin(), axes.end());
    uint32_t axesDataNum = 1;   
    shapeDimNum = context->GetInputShape(0)->GetStorageShape().GetDimNum();
    OP_CHECK_IF(shapeDimNum > 20, OP_LOGE(context, "dimNum is too much"), return ge::GRAPH_FAILED); 
    shapeNum.resize(shapeDimNum, 0);
    strides.resize(shapeDimNum, 0);
    for (uint32_t i = 0; i < shapeDimNum; i++) {
        shapeNum[i] = context->GetInputShape(0)->GetStorageShape().GetDim(i);
        OP_CHECK_IF(shapeNum[i] <= 0, OP_LOGE(context, "shapeNum is 0"), return ge::GRAPH_FAILED);
        int32_t totalStride = inputNum;
        for (uint32_t j = 0; j < i + 1; j++) {
            totalStride /= shapeNum[j];
        }
        strides[i] = totalStride;
    }
    for (uint32_t i = 0; i < axesNum; i++) {
        if (axes[i] < 0) {
            axes[i] = axes[i] + shapeDimNum;
        }
        axesDataNum *= shapeNum[axes[i]];
    }
    
    reduceMask.resize(shapeDimNum, false);
    for (uint32_t i = 0; i < axesNum; i++) {
        reduceMask[axes[i]] = true;
    }
    
    //core
    coreDataNum = inputNum / axesDataNum / coreNum;
    tailCoreNum = inputNum / axesDataNum % coreNum;  
    //cycle
    cycleDataNum = shapeNum[axes[axesNum - 1]];
    cycleNum = axesDataNum / cycleDataNum;
    //tile
    uint32_t align32Length = ((cycleDataNum * inputTypeLen + BLOCK_SIZE - 1) / BLOCK_SIZE) * BLOCK_SIZE;
    uint32_t inputBlockNum = align32Length / BLOCK_SIZE;
    uint32_t ubDataNumber = (inputTypeLen == 2) ? UB_DATA_NUM_FLOAT16 : UB_DATA_NUM_FLOAT32; 
    uint32_t tileBlockNum = (ubSize / BLOCK_SIZE / BUFFER_NUM) / ubDataNumber;
    tileDataNum = tileBlockNum * BLOCK_SIZE / 32; //DataCopyPad，GatherMask的32位对齐
    uint32_t tileNum = inputBlockNum / tileBlockNum;
    finaltileNum = (inputBlockNum % tileBlockNum == 0) ? tileNum : tileNum + 1;
    tailDataNum = cycleDataNum - (tileDataNum * tileNum);
    return ge::GRAPH_SUCCESS;
}

static ge::graphStatus ReduceLogSumExpTilingFunc(gert::TilingContext* context)
{
    ReduceLogSumExpTilingData* tiling = context->GetTilingData<ReduceLogSumExpTilingData>();
    OP_CHECK_NULL_WITH_CONTEXT(context, tiling);
    OP_CHECK_IF(
        memset_s(tiling, sizeof(ReduceLogSumExpTilingData), 0, sizeof(ReduceLogSumExpTilingData)) != EOK,
        OP_LOGE(context, "set tiling data error"), return ge::GRAPH_FAILED);
    // 获取平台运行信息
    uint64_t ubSize;
    uint64_t coreNum;
    uint32_t inputTypeLen;
    uint32_t axesNum, shapeDimNum;
    vector<int32_t> axes, shapeNum, strides;
    vector<bool> reduceMask;
    ge::graphStatus ret = GetPlatformInfo(context, ubSize, coreNum, inputTypeLen);
    OP_CHECK_IF(ret != ge::GRAPH_SUCCESS, OP_LOGE(context, "GetPlatformInfo error"), return ge::GRAPH_FAILED);
    // 获取输入数据信息
    uint32_t inputNum, tileDataNum, finaltileNum, tailDataNum, coreDataNum, tailCoreNum, cycleDataNum, cycleNum;
    ret = GetTilingAttrsInfo(context, ubSize, inputNum, tileDataNum, finaltileNum, tailDataNum,
                           coreDataNum, tailCoreNum, cycleDataNum, cycleNum, axesNum, shapeDimNum, 
                           axes, shapeNum, strides, coreNum, inputTypeLen, reduceMask);
    OP_CHECK_IF(ret != ge::GRAPH_SUCCESS, OP_LOGE(context, "GetTilingAttrsInfo error"), return ge::GRAPH_FAILED);

    // 设置tiling数据
    tiling->cycleDataNum = static_cast<uint32_t>(cycleDataNum);
    tiling->cycleNum = static_cast<uint32_t>(cycleNum);  
    tiling->coreDataNum = static_cast<uint32_t>(coreDataNum);
    tiling->tailCoreNum = static_cast<uint32_t>(tailCoreNum);   
    tiling->tileDataNum = static_cast<uint32_t>(tileDataNum);
    tiling->tileNum = static_cast<uint32_t>(finaltileNum);
    tiling->tailDataNum = static_cast<uint32_t>(tailDataNum);   
    tiling->inputNum = static_cast<uint32_t>(inputNum);
    tiling->axesNum = static_cast<uint32_t>(axesNum);   
    tiling->shapeDimNum = static_cast<uint32_t>(shapeDimNum);
    
    std::copy(axes.begin(), axes.end(), tiling->axes);
    std::copy(shapeNum.begin(), shapeNum.end(), tiling->shapeNum);
    std::copy(strides.begin(), strides.end(), tiling->strides);
    std::copy(reduceMask.begin(), reduceMask.end(), tiling->reduceMask);
    
    // 计算workspace大小
    OP_CHECK_IF(GetWorkspaceSize(context) != ge::GRAPH_SUCCESS, OP_LOGE(context, "GetWorkspaceSize error"), return ge::GRAPH_FAILED);
    uint32_t tilingKey = 0;
    tilingKey = GET_TPL_TILING_KEY(0);
    context->SetTilingKey(tilingKey);
    context->SetBlockDim(coreDataNum == 0 ? tailCoreNum: coreNum);
    
    return ge::GRAPH_SUCCESS;
}

// tiling注册入口.
IMPL_OP_OPTILING(ReduceLogSumExp).Tiling(ReduceLogSumExpTilingFunc).TilingParse<ReduceLogSumExpCompileInfo>(TilingParseForReduceLogSumExp);
} // namespace optiling